#include "TcpConnection.h"

using namespace clusnet::net;

TcpConnection::~TcpConnection()
{
    assert(state_ == ConnectionState::CLOSED);
    if (!sendMsgs_.empty())
    {
        LOG(ERROR) << sendMsgs_.size() << " messages will be discarded from " << localEndpoint().address() << ":" << localEndpoint().port() << " to " << remoteEndpoint_.address() << ":" << remoteEndpoint_.port();
    }
}

void TcpConnection::initConnect(bool nodelay, bool keepAlive, const TcpWriteCompleteCallback& writeCallback, const TcpMessageCallback& messageCallback, const TcpCloseCallback& closeCallback)
{
    setNoDelay(nodelay);
    setKeepAlive(keepAlive);
    setRemoteEndpoint();
    setWriteCompleteCallback(writeCallback);
    setMessageCallback(messageCallback);
    setCloseCallback(closeCallback);

    state_ = ConnectionState::CONNECTED;
}

void TcpConnection::safeShutDown()
{
    auto self(shared_from_this());
    boost::asio::dispatch(io_, boost::asio::bind_executor(strand_, [this, self]()
    {
        if (state_ != ConnectionState::CONNECTED)
        {
            return;
        }
        state_ = ConnectionState::SHUTDOWNED;
        socket_.shutdownWrite();
        closeTimer_.expires_after(std::chrono::seconds(3));
        closeTimer_.async_wait(boost::asio::bind_executor(strand_, [this, self](const boost::system::error_code& ec) -> void
        {
            if (!ec)
            {
                socket_.close();
            }
        }));
    }));
}

void TcpConnection::close()
{
    closeTimer_.cancel();
    assert(closeCallback_ != nullptr);
    auto self = shared_from_this();
    closeCallback_(self);
    state_ = ConnectionState::CLOSED;
}

void TcpConnection::safeAsyncWrite(TcpDataPacketPtr const& data)
{
    data->encodeHeader();
    auto self(shared_from_this());
    boost::asio::dispatch(io_, boost::asio::bind_executor(strand_, [this, self, data]()
    {
        bool write_in_progress = !sendMsgs_.empty();
        sendMsgs_.emplace_back(data);
        if (!write_in_progress)
        {
            asyncWrite();
        }
    }));
}

void TcpConnection::safeStartReceive()
{
    auto self(shared_from_this());
    boost::asio::dispatch(io_, boost::asio::bind_executor(strand_, [this, self]()
    {
        receiveHeader();
    }));
}

void TcpConnection::asyncWrite()
{
    auto self(shared_from_this());
    boost::asio::async_write(socket_.socket(), boost::asio::buffer(sendMsgs_.front()->cdata(), sendMsgs_.front()->dataLength()), boost::asio::bind_executor(strand_, [this, self](boost::system::error_code const& ec, std::size_t length)
    {
        if (ec)
        {
            safeShutDown();
            return;
        }
        if (length != sendMsgs_.front()->dataLength())
        {
            LOG(WARNING) << "async_write failed. data length not equal. expected send len: " << sendMsgs_.front()->dataLength() << " actual send len: " << length;
            safeShutDown();
            return;
        }

        if (writeCallback_)
        {
            auto self(shared_from_this());
            writeCallback_(self, sendMsgs_.front());
        }

        sendMsgs_.pop_front();
        if (!sendMsgs_.empty())
        {
            asyncWrite();
        }
    }));
}

void TcpConnection::receiveHeader()
{
    auto self(shared_from_this());
    recvMsg_.reset(new TcpDataPacket);
    boost::asio::async_read(socket_.socket(), boost::asio::buffer(recvMsg_->header(), recvMsg_->headerLength()), boost::asio::bind_executor(strand_, [this, self](boost::system::error_code const& ec, std::size_t length)
    {
        if (ec)
        {
            close();
            return;
        }

        if (length != recvMsg_->headerLength())
        {
            close();
            return;
        }
        receiveBody();
    }));
}

void TcpConnection::receiveBody()
{
    recvMsg_->decodeHeader();
    auto self(shared_from_this());
    boost::asio::async_read(socket_.socket(), boost::asio::buffer(recvMsg_->body(), recvMsg_->bodyLength()), boost::asio::bind_executor(strand_, [this, self](boost::system::error_code const& ec, std::size_t length)
    {
        if (ec)
        {
            close();
            return;
        }

        if (length != recvMsg_->bodyLength())
        {
            close();
            return;
        }
        assert(messageCallback_ != nullptr);
        auto self = shared_from_this();
        messageCallback_(self, std::move(recvMsg_)); 

        receiveHeader();
    }));
}
